from langchain_ollama import ChatOllama
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

# 初始化LLM
llm = ChatOllama(
    model="deepseek-r1:8b",
    base_url="http://localhost:11434",  # Ollama 服务地址
    temperature=0.7,  # 创造性程度
    num_predict=1024,  # 最大生成长度
    options={'stop': ['</think>', '</think>']} # 遇到这些标记就停止
)

# 定义不同的专业角色模板
flower_care_template = """你是一个经验丰富的园丁，擅长解答关于养花育花的问题。{input}"""
flower_deco_template = """你是一位网红插花大师，擅长解答关于鲜花装饰的问题。{input}"""
default_template = """你是一个通用助手，回答各类问题。{input}"""

# 创建不同的prompt
flower_care_prompt = ChatPromptTemplate.from_template(flower_care_template)
flower_deco_prompt = ChatPromptTemplate.from_template(flower_deco_template)
default_prompt = ChatPromptTemplate.from_template(default_template)

# 创建不同的chain
flower_care_chain = flower_care_prompt | llm | StrOutputParser()
flower_deco_chain = flower_deco_prompt | llm | StrOutputParser()
default_chain = default_prompt | llm | StrOutputParser()

# 创建路由决策prompt
router_template = """你需要分析用户的问题，并决定使用哪个专业角色来回答。

可用的专业角色：
1. flower_care: 适合回答关于鲜花护理的问题
2. flower_decoration: 适合回答关于鲜花装饰的问题

如果问题明确属于某个专业领域，请返回对应的角色名称。
如果无法确定或问题超出上述范围，请返回'default'。

请仅返回角色名称，不要添加任何其他文本。

用户问题：{input}"""

router_prompt = ChatPromptTemplate.from_template(router_template)
router_chain = router_prompt | llm | StrOutputParser()

# 舍弃router_chain，直接根据问题判断路由，进入响应的chain，否则进入default_chain

def route_query(query):
    # 打印路由提示信息（类似于原代码的打印）
    print("路由提示:\n", router_template.format(input="{input}"))
    
    # 获取路由决策
    route = router_chain.invoke({"input": query}).strip().lower()
    
    # 确保路由结果有效
    if "flower_care" in route:
        selected_chain = flower_care_chain
    elif "flower_decoration" in route:
        selected_chain = flower_deco_chain
    else:
        selected_chain = default_chain
    
    # 执行选择的chain
    return selected_chain.invoke({"input": query})

result = route_query("如何给玫瑰花浇水")
print("\n回答结果:\n", result)